import os
import sys
# from data_augmenter import dataAugmenter
# from pcb import pcb
# from graph import graph
# from .geometry import BBox, Point, Tx, Vector
# from .agent import agent as agent
# from .agentP.parameters import parameters as agent_parameters
# from core.environment.tracker import tracker
# from pcbDraw import draw_board_from_board_and_graph_with_debug, draw_ratsnest_with_board
import numpy as np
from datetime import datetime
import uuid
import torch

from envs.sch.Part import Part
from .sch.schDraw import create_box_and_random_placement
from skidl.schematics.place import get_enclosing_bbox
from skidl.schematics.debug_draw_rl import (
    draw_end,
    draw_pause,
    draw_placement,
    draw_redraw,
    draw_start,
    draw_text,
)
from envs.sch.gen_node import gen_node4Env
from PIL import Image
class SchEnv:
    # This method is called when an object is created.
    # It's purpose is to initialize the object.
    def __init__(self,
                 max_episode_steps=1000,
                 if_save_pic=True,
                 if_render_pic=True,
                 ):

        '''来自skidl的原理图原始数据'''
        self.node = gen_node4Env()
        self.parts = self.node.parts
        self.nets = self.node.parts[0].circuit.nets
        self.n_agents = len(self.parts)

        '''画图控制'''
        self.if_sava_pic = if_save_pic
        self.if_render_pic = if_render_pic

        '''the necessary env information when you design a custom env'''
        self.env_name = "SchEnv"  # the name of this env.
        self.num_envs = 1  # the number of sub env is greater than 1 in vectorized env.
        self.max_step = max_episode_steps  # the max step number of an episode.
        '''生成rlParts'''
        self.rlParts = []
        self.initialize_environment_state_from_sch()
        self.state_dim = self.rlParts[0].observation_space.shape[0]  # feature number of state
        self.action_dim = self.rlParts[0].action_space.shape[0]  # feature number of action
        self.if_discrete = False  # discrete action or continuous action
        self.uid = uuid.uuid1().hex
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        '''states, rewards, dones, info'''
        self.states = torch.zeros(( self.n_agents, self.state_dim), dtype=torch.float32).to(self.device)
        # self.actions = torch.zeros(( self.n_agents, self.action_dim), dtype=torch.float32).to(self.device)
        self.rewards = torch.zeros(self.n_agents, dtype=torch.float32).to(self.device)
        self.dones = torch.zeros(self.n_agents, dtype=torch.bool).to(self.device)

    def reset(self):

        for i in range(self.n_agents):
            self.states[i] = self.rlParts[i].reset()
        return self.states


    def step(self,actions):
        if self.if_sava_pic:
            # Draw the placement for debug purposes.
            bbox = get_enclosing_bbox(self.parts)
            draw_scr, draw_tx, draw_font = draw_start(bbox)

        for i in range(self.n_agents):
            (self.states[i],
             self.rewards[i],
             self.dones[i],
             _ )= self.rlParts[i].step(actions[i])

        if self.if_sava_pic:
            draw_placement(self.parts, self.nets, draw_scr, draw_tx, draw_font)
            path = os.path.join("./", datetime.now().strftime('%Y-%m-%d_%H-%M-%S'))+self.uid+".png"
            draw_pause(draw_scr,path)
            img_pil = Image.open(path)
            # 读取图像并转换为 PyTorch Tensor
            # 注意：TensorBoard 需要的是 CHW 格式的 Tensor，即 (C, H, W)，并且数据类型应该是 float32，值在 [0, 1] 之间
            img_tensor = np.array(img_pil)
            # print(img_tensor.shape)
            # print(type(img_tensor))

        return self.states, self.rewards, self.dones, None



    def initialize_environment_state_from_sch(self, **options):
        '''
        初始化原理图
        :param options:
        :return:
        '''
        self.bboxWidth, self.bboxHeight = create_box_and_random_placement(self.parts, **options)
        self.agents = []
        for i in range(self.n_agents):
            rlPart = Part(self.parts[i],
                          self.max_step,
                          self.bboxHeight,
                          self.bboxWidth,
                          self
                          )
            self.rlParts.append(rlPart)


if __name__ == "__main__":
    print(gen_node4Env())